# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import json
import time
import asyncio
from decimal import Decimal
from typing import List, Dict
from fastapi import Request
from tortoise.models import in_transaction
from exception import AppException
from common.utils.urls import UrlUtil
from common.utils.tools import ToolsUtil
from common.utils.files import FilesUtil
from common.utils.config import ConfigUtil
from common.enums.wallet import WalletEnum
from common.models.users import UserModel
from common.models.users import UserWalletModel
from common.models.dev import DevModelsModel
from common.models.docs import DocsFilesModel
from common.models.docs import DocsPromptModel
from common.models.docs import DocsRecordsModel
from common.cache.pool_cache import KeyPoolCache
from common.service.retrieval_logic import RecallDocsLogic, RecallUtil
from common.chain.chain_server import FlowsSchema, ChatUtils, AIChannelFactory


class ChatDocsWidget:
    def __init__(self, request: Request, post: dict):
        self.start_time = time.time()
        self.request = request
        self.user_id: int = int(post.get("user_id"))
        self.ip: str = str(post.get("ip"))
        self.type: str = str(post.get("type"))  # [abstract, translate, dialogue]
        self.code: str = str(post.get("code"))
        self.question: str = str(post.get("question"))
        self.instruct: str = ""

        # 指令处理
        if self.question.startswith("@"):
            array = self.question.split(" ")
            if len(array) >= 2:
                self.instruct = array.pop(0).strip()
                self.question = " ".join(array).strip()
        print(self.question)
        print(self.instruct)
        # 实例对象
        self.user: UserModel
        self.docFile: DocsFilesModel
        self.docMode: DocsPromptModel

        # 文档内容
        self.quotes: list = []
        self.document: str = ""
        self.pageList: List[dict] = []

        # 收费方式
        self.charge_method = 0  # [0=免费,1=模型,2=固定]
        self.charge_amount = 0  # [固定费用]

        # 对话模型
        self.chat_mark: str = ""
        self.chat_channel: str = ""
        self.chat_alias: str = ""
        self.chat_model: str = ""
        self.chat_price: int = 0
        self.chat_config: Dict[str, any] = {}

        # 向量模型
        self.vector_mark: str = ""
        self.vector_alias: str = ""
        self.vector_model: str = ""
        self.vector_price: int = 0

        # 检索配置
        self.context_num: int = 3
        self.search_mode: str = "similar"
        self.search_tokens: int = 5
        self.search_similar: Decimal = Decimal(0.4)
        self.ranking_status: bool = False
        self.ranking_models: str = ""
        self.ranking_score: Decimal = Decimal(0.5)

        # 响应数据
        self.think: str = ""  # 思考内容
        self.reply: str = ""  # 答复内容
        self.reply_dict: Dict[int, str] = {}
        self.flows_usage: List[FlowsSchema] = []

    async def chat(self):
        try:
            # 初始化数据
            await self._initialize()

            # 按类型处理
            messages = []
            if self.type == "abstract":
                await self._load_model(self.docMode.models_abstract)
                messages.append([
                    {"role": "system", "content": self.docMode.prompt_abstract},
                    {"role": "user", "content": json.dumps({
                        "filename": self.docFile.name,
                        "document": self.document
                    }, ensure_ascii=False)}
                ])
            elif self.type == "translate":
                await self._load_model(self.docMode.models_translate)
                for key, item in enumerate(self.pageList):
                    messages.append([
                        {"role": "system", "content": self.docMode.prompt_translate},
                        {"role": "user", "content": json.dumps({
                            "index": key + 1,
                            "page_no": item["page_no"],
                            "filename": self.docFile.name,
                            "language": self.question,
                            "document": item["content"]
                        }, ensure_ascii=False)}
                    ])
            elif self.type == "dialogue":
                await self._load_model(self.docMode.models_dialogue)
                if self.instruct in ["@全文", "@fullText", "@full"]:
                    document = self.document
                else:
                    document = await self._call_learning()

                history = await self._get_history()
                messages.append([
                    {"role": "system", "content": self.docMode.prompt_dialogue},
                    *history,
                    {"role": "user", "content": json.dumps({
                        "question": self.question,
                        "filename": self.docFile.name,
                        "document": document
                    }, ensure_ascii=False)}
                ])

            # 发起流式提问
            usages: dict = {}
            cancel: bool = False
            start_time = time.time()
            for index, context in enumerate(messages):
                # SSE迭代输出
                try:
                    async for output in self._send_sse(index, context):
                        yield output
                except asyncio.CancelledError:
                    cancel = True

                # 计算Tokens
                answer: str = self.reply_dict[index]
                usage: dict = ChatUtils.usage(context, answer)
                usages = {
                    "total_tokens": usage["total_tokens"] + usages.get("total_tokens", 0),
                    "prompt_tokens": usage["prompt_tokens"] + usages.get("prompt_tokens", 0),
                    "completion_tokens": usage["completion_tokens"] + usages.get("completion_tokens", 0)
                }

                # 客户端中断
                if cancel:
                    break

            # 对话消耗情况
            if usages:
                consume_points = ChatUtils.compute_price(usages["total_tokens"], self.chat_price)
                self.flows_usage.append(FlowsSchema(
                    name="AI 对话",
                    scene=str("chat"),
                    alias=self.chat_alias,
                    model=self.chat_model,
                    task_time=f"{(time.time() - start_time):.2f}",
                    use_points=consume_points,
                    total_tokens=usages["total_tokens"],
                    prompt_tokens=usages["prompt_tokens"],
                    completion_tokens=usages["completion_tokens"],
                ))

            # 客户主动取消
            if cancel:
                task = asyncio.create_task(self._save_record())
                await asyncio.shield(task)
            else:
                await self._save_record()
                yield ChatUtils.output_success(idx="", model=self.chat_alias)
                yield "data: [DONE]"
        except Exception as e:
            yield ChatUtils.output_error(model=self.chat_model, message=str(e))

    async def _send_sse(self, index: int, messages: List[dict]):
        """ 发送SSE """
        apikey: dict = await self._get_apikey()
        aiServer = AIChannelFactory.create(self.chat_channel, self.chat_config, apikey)
        stream = await aiServer.llm(messages)
        async for chunk in stream:
            if not chunk.choices:
                raise Exception(f"{chunk.code}: {chunk.message}")

            think = ""
            event = "chat"
            delta = chunk.choices[0].delta.model_dump()

            if chunk.choices[0].finish_reason == "normal":
                event = "finish"
            if delta.get("reasoning_content"):
                think = chunk.choices[0].delta.reasoning_content
                self.think += chunk.choices[0].delta.reasoning_content
            if delta.get(f"content"):
                self.reply += chunk.choices[0].delta.content

            answer = self.reply_dict.get(index, "")
            self.reply_dict[index] = answer + chunk.choices[0].delta.content

            yield ChatUtils.output_success(
                idx=chunk.id,
                event=event,
                model=self.chat_model,
                index=chunk.choices[0].index,
                data=chunk.choices[0].delta.content,
                think=think
            )

    async def _initialize(self):
        """ 初始化数据 """
        # 查验文件
        self.docFile = await DocsFilesModel.filter(code=self.code, user_id=self.user_id).first()
        if self.docFile is None:
            raise Exception("文档文件不存在")
        if self.docFile.is_delete:
            raise Exception("文档文件已被删")

        # 查验模式
        self.docMode = await DocsPromptModel.filter(code=self.docFile.mode, is_delete=0).first()
        if self.docMode is None:
            raise Exception("选择的对话模式不存在")
        if self.docMode.is_disable:
            raise Exception("选择的对话模式已下架")

        # 查询配置
        config = await ConfigUtil.get("docs") or {}
        if not config.get("status"):
            raise AppException("阅读功能已暂停使用")

        # 收费方式: [0=免费, 1=模型, 2=固定]
        self.charge_method = int(config.get("charge_method", "0"))
        self.charge_amount = int(config.get("charge_amount", "0"))

        # 检索条件
        self.search_mode = str(config.get("search_mode", "similar"))
        self.search_tokens = int(config.get("search_tokens", "5000"))
        self.search_similar = Decimal(config.get("search_similar", "0.4"))
        self.ranking_status = bool(int(config.get("ranking_status", "0")))
        self.ranking_models = str(config.get("ranking_models", ""))
        self.ranking_score = Decimal(config.get("ranking_score", "0.5"))

        # 向量模型
        emb_model: str = config.get("vector_models", "")
        vector_model = await DevModelsModel.check_models(emb_model, "vector")
        self.vector_mark = vector_model["mark"]
        self.vector_alias = vector_model["alias"]
        self.vector_model = vector_model["model"]
        self.vector_price = vector_model["price"]

        # 查询用户
        self.user = await UserModel.filter(id=self.user_id).first()
        if not self.user:
            raise Exception("用户身份异常")

        # 验证余额
        if self.charge_method > 0 and self.user.points <= int(0):
            error = "账户余额不足,请充值后重试!"
            if self.chat_price:
                raise AppException(error)
            if self.vector_price:
                raise AppException(error)
            if self.ranking_status:
                raise AppException(error)

        # 读取文件
        await self._read_files()

    async def _load_model(self, llm: str):
        """ 加载模型 """
        config: dict = json.loads(llm or "{}")
        models: str = config.get("model", "")
        chat_model = await DevModelsModel.check_models(models, "chat")
        self.chat_mark = chat_model["mark"]
        self.chat_channel = chat_model["channel"]
        self.chat_alias = chat_model["alias"]
        self.chat_model = chat_model["model"]
        self.chat_price = int(chat_model["price"])
        self.chat_config = chat_model["config"]
        for key, val in self.chat_config.items():
            if config.get(key) is not None:
                self.chat_config[key] = config.get(key)
        self.chat_config["model"] = self.chat_model

    async def _get_apikey(self):
        """ 获取密钥 """
        apikey = {}
        if self.chat_config["check_key"]:
            apikey = await KeyPoolCache(scene="chat", channel=self.chat_mark).get_key() or {}
            if not apikey:
                raise AppException("当前模型尚未配置密钥,请与管理员联系")
        return apikey

    async def _get_history(self):
        """ 历史对话 """
        messages = []
        if self.context_num > 0:
            records = await (DocsRecordsModel
                             .filter(user_id=self.user_id)
                             .filter(file_id=self.docFile.id)
                             .filter(scene=self.type)
                             .filter(is_clear=0)
                             .filter(is_delete=0)
                             .order_by("-id")
                             .limit(self.context_num)
                             .all())
            records = list(reversed(records))
            for record in records:
                messages.append({"role": "user", "content": record.question})
                messages.append({"role": "assistant", "content": record.answer})
        return messages

    async def _call_learning(self):
        # 问题优化
        start_time = time.time()
        questions = [self.question]

        # 向量模型
        vector_model = f"{self.vector_mark}:{self.vector_model}"

        # 发起检索
        results = []
        for query in questions:
            # 按模式检索
            if self.search_mode == "similar":
                lists = await RecallDocsLogic.embedding_recall(
                    model=vector_model,
                    query=query,
                    file_id=self.docFile.id,
                    limit=100
                )
                results.append({"k": 60, "list": lists})
            elif self.search_mode == "full":
                lists = await RecallDocsLogic.full_text_recall(
                    query=query,
                    file_id=self.docFile.id,
                    limit=100
                )
                results.append({"k": 60, "list": lists})
            elif self.search_mode == "mix":
                lists = await RecallDocsLogic.mixed_recall(
                    model=vector_model,
                    query=query,
                    file_id=self.docFile.id,
                )
                results.append({"k": 60, "list": lists})

        # 检索费用
        if self.search_mode in ["similar", "mix"]:
            use_content: str = "".join(questions)
            total_tokens: int = ChatUtils.compute_tokens(use_content)
            consume_points: int = ChatUtils.compute_price(total_tokens, self.vector_price)
            # 流程登记
            self.flows_usage.append(FlowsSchema(
                name="文档搜索",
                scene="vector",
                alias=self.vector_alias,
                model=self.vector_model,
                task_time=f"{(time.time() - start_time):.2f}",
                use_points=consume_points,
                total_tokens=total_tokens,
                prompt_tokens=total_tokens,
                completion_tokens=0
            ))

        # 结果处理(1): RRF重排
        results = RecallUtil.rrf_concat_results(results)

        # 结果处理(2): 重排模型
        similar: Decimal = self.search_similar
        if self.ranking_status:
            models = self.ranking_models
            similar = self.ranking_score
            data = await RecallUtil.re_ranker_results(results, question=self.question, models=models)
            self.ranking_status = bool(data["status"])
            if self.ranking_status:
                results = data["results"]
                self.flows_usage.append(data["usage"])

        # 结果处理(3): 相似度过滤
        results = RecallUtil.filter_max_score(
            results=results,
            mode=self.search_mode,
            rerank=self.ranking_status,
            similar=similar
        )

        # 结果处理(4): MaxTokens过滤
        results = RecallUtil.filter_max_tokens(
            results=results,
            max_tokens=self.search_tokens
        )

        # 结果处理(5): 整合引用知识
        index: int = 0
        quotes: list = []
        content: str = ""
        for item in results:
            quotes.append(item)

            if index > 0:
                content += "\n------------\n\n"

            index += 1
            content += item["chunk_texts"] + "\n"

        # 返回检索结果
        self.quotes = quotes
        self.document = content
        return self.document

    async def _save_record(self):
        """ 保存记录 """
        async with in_transaction("mysql"):
            # 重查用户
            self.user = await UserModel.filter(id=self.user_id).first()

            # 消耗的Tokens
            total_consume_tokens: int = 0
            for item in self.flows_usage:
                total_consume_tokens += item.total_tokens

            # 消耗的积分数
            if self.charge_method == 0:  # 免费
                total_consume_points: int = 0
            elif self.charge_method == 2:  # 固定
                total_consume_points: int = self.charge_amount
            else:  # 模型
                total_consume_points: int = 0
                for item in self.flows_usage:
                    total_consume_points += item.use_points

            # 用户扣费
            if total_consume_points > 0:
                surplus_balance: int = max(self.user.points - total_consume_points, 0)
                await UserModel.filter(id=self.user_id).update(points=surplus_balance)

            # 更新文件
            self.docFile = await DocsFilesModel.filter(id=self.docFile.id).first()
            await DocsFilesModel.filter(id=self.docFile.id).update(
                total_chat_count=self.docFile.total_chat_count + 1,
                total_use_points=self.docFile.total_use_points + total_consume_points,
                total_use_tokens=self.docFile.total_use_tokens + total_consume_tokens
            )

            # 写入记录
            salt: str = self.ip + str(self.user_id) + str(time.time())
            code: str = ToolsUtil.make_md5_str(salt + str(self.docFile.id))
            flows_usage = [flow.model_dump() for flow in self.flows_usage]
            await DocsRecordsModel.create(
                user_id=self.user_id,
                file_id=self.docFile.id,
                code=code,
                scene=self.type,
                question=self.question,
                answer=self.reply,
                use_tokens=total_consume_tokens,
                use_points=total_consume_points,
                ip=self.ip,
                quotes=json.dumps(self.quotes, ensure_ascii=False),
                flows_usage=json.dumps(flows_usage, ensure_ascii=False),
                task_time=f"{time.time() - self.start_time:.2f}",
                create_time=int(time.time()),
                update_time=int(time.time())
            )

            # 写入日志
            if total_consume_points > 0:
                projects = {"abstract": "文档总结", "dialogue": "文档对话", "translate": "文档翻译"}
                await UserWalletModel.dec(
                    scene=WalletEnum.POINTS,
                    user_id=self.user_id,
                    change_type=WalletEnum.POINTS_DEC_DOCS_CHAT,
                    change_amount=total_consume_points,
                    remarks=WalletEnum.get_change_type_msg(WalletEnum.POINTS_DEC_DOCS_CHAT),
                    project=projects.get(self.type),
                    source_sn=self.docFile.code,
                    source_id=self.docFile.id,
                    additional=json.dumps(flows_usage, ensure_ascii=False)
                )

    async def _read_files(self):
        """ 读取文件 """
        # 文档路径
        file_path = UrlUtil.to_runtime_path(self.docFile.txt_path)

        # 文档不存在: 重新解析文档
        if not os.path.exists(file_path):
            # 本地PDF不存在则远程下载
            pdf_path: str = UrlUtil.to_root_path(self.docFile.pdf_path)
            if not os.path.exists(pdf_path):
                url: str = await UrlUtil.to_absolute_url(self.docFile.pdf_path)
                await ToolsUtil.download_file(url=url, save_path=pdf_path)

            # 重新读取PDF的文档内容
            pages = await FilesUtil.parse_pdf(pdf_path)
            texts = json.dumps(pages, indent=4, ensure_ascii=False)
            await ToolsUtil.write_file(file_path, texts)

        # 文档存在的: 读取文档内容
        try:
            texts = ""
            pages = await ToolsUtil.read_file(file_path)
            self.pageList = json.loads(pages)
            for page in self.pageList:
                texts += page.get("content", "")
            self.document = texts
            return self.document
        except Exception as e:
            print(str(e))
